# mesh.py
# Austin Abrams, Washington University in St. Louis
# 1/8/2010
	
# Writes a .dae (COLLADA format) polygon file

from models import Correspondence
import math
from datetime import datetime
from matfunc import *
from helperFunctions import *

# the higher this number is, the more we sample the texture, but the more space the file takes.
# 3 seems to give a good enough texture.
SUBDIVISION_K = 3

def writeMeshFile(correspondences, triangulation, width, height):
	assert len(triangulation) % 3 == 0
	noTriangles = len(triangulation) / 3

	# Step 1:  normalize uv coordinates.
	u = []
	v = []
	for i in correspondences:
		u.append(float(i.u) / width);
		v.append(1 - float(i.v) / height);
	
	# Step 2:  translate (lat,lon,alt) into meters, centered around point (0,0,0).
	origin = correspondences[0]
	originLat = origin.latitude
	originLon = origin.longitude
	originAlt = origin.altitude
	x = []
	y = []
	z = []
	metersPerLongitude = distance(originLat, originLon, originLat, originLon + 0.00001) / 0.00001;
	metersPerLatitude  = distance(originLat, originLon, originLat + 0.00001, originLon) / 0.00001;
	for i in correspondences:
		x.append((i.longitude - originLon) * metersPerLongitude)
		y.append((i.latitude - originLat)  * metersPerLatitude)
		z.append(i.altitude - originAlt)
		
	# subdivide the quadrilateral.
	if (len(x) == 4):
		x,y,z,u,v,nx,ny,nz,triangles = subdivideQuadrilateralHomography(x,y,z,u,v,SUBDIVISION_K)
		#x,y,z,u,v,nx,ny,nz,triangles = subdivideQuadrilateralInterp(x,y,z,u,v,SUBDIVISION_K)
	
	else:
		# Step 3: build triangle spec.
		triangles = []
		for i in range(0, len(triangulation), 3):
			tri = [ [], [], [] ]
			tri[0] = {'pos': triangulation[i],   'tex': triangulation[i],   'nor': i/3};
			tri[1] = {'pos': triangulation[i+1], 'tex': triangulation[i+1], 'nor': i/3};
			tri[2] = {'pos': triangulation[i+2], 'tex': triangulation[i+2], 'nor': i/3};
			triangles.append(tri)
		
		# find the flat-shading normals.
		nx,ny,nz = findNormals(x,y,z,triangles)
	
	# Step 4: build it!
	daePositions, daeNormals, daeUVs, daeTriangles, nPositions, nNormals, nUVs, nTriangles = writeDAE(x,y,z,u,v,nx,ny,nz,triangles)
	
	return originLat, originLon, originAlt, daePositions, daeNormals, daeUVs, daeTriangles, nPositions, nNormals, nUVs, nTriangles
	
	
def findNormals(x,y,z,triangles):
	nx,ny,nz = [], [], []
	for tri in triangles:
		i0 = tri[0]['pos']
		i1 = tri[1]['pos']
		i2 = tri[2]['pos']
		x0 = x[i0]
		y0 = y[i0]
		z0 = z[i0]
		x1 = x[i1]
		y1 = y[i1]
		z1 = z[i1]
		x2 = x[i2]
		y2 = y[i2]
		z2 = z[i2]
		n = normal(x0,y0,z0,x1,y1,z1,x2,y2,z2)
		nx.append(n[0])
		ny.append(n[1])
		nz.append(n[2])
	return nx,ny,nz
	

# given:
#    x,y, and z are a list of positions,
#    u and v are image coordinates in [0,1]
#    nx,ny,nz are a list of normals
#    triangles is a dictionary, where
#         triangles[i][0], triangles[i][1], and triangles[i][2] are individual vertices, where
#             triangles[i][j]['pos'] indexes into (x,y,z)
#             triangles[i][j]['tex'] indexes into (u,v)
#             triangles[i][j]['nor'] indexes into (nx,ny,nz).
#    if doubleSided is set to true, then for each triangle with normal vector N,
#    a second triangle is made with normal vector -N.
# this returns:
#    daePositions, daeNormals, daeUVs, daeTriangles,
#    nPositions, nNormals, nUVs, nTriangles
#    the arrays that, when placed in the right locations, makes a COLLADA file.
def writeDAE(x,y,z,u,v,nx,ny,nz,triangles, doubleSided=True):
	assert len(x) == len(y)
	assert len(y) == len(z)
	
	assert len(u) == len(v)
	
	assert len(nx) == len(ny)
	assert len(ny) == len(nz)
	
	# POSITIONS
	daePositions = ''
	nPositions = 0
	for i in range(len(x)):
		daePositions += '%1.4f %1.4f %1.4f ' % (x[i], y[i], z[i])
		nPositions += 3
			
	if doubleSided:
		for i in range(len(x)):
			# find a normal that uses this point.
			normalI = None
			for tri in triangles:
				if normalI:
					break
				for point in tri:
					if point['pos'] == i:
						normalI = point['nor']
						break
			altNx = -nx[normalI] * 0.01
			altNy = -ny[normalI] * 0.01
			altNz = -nz[normalI] * 0.01
			daePositions += '%1.4f %1.4f %1.4f ' % (x[i]+altNx, y[i]+altNy, z[i]+altNz)
			nPositions += 3

	# NORMALS
	daeNormals = ''
	nNormals = 0
	for i in range(len(nx)):
		daeNormals += '%1.4f %1.4f %1.4f ' % (nx[i], ny[i], nz[i])
		nNormals += 3
	
	if doubleSided:
		for i in range(len(nx)):
			daeNormals += '%1.4f %1.4f %1.4f ' % (-nx[i], -ny[i], -nz[i])
			nNormals += 3

	# UVs
	daeUVs = ''
	nUVs = 0
	for i in range(len(u)):
		daeUVs += '%1.4f %1.4f ' % (u[i], v[i])
		nUVs += 2

	# TRIANGLES
	daeTriangles = ''
	nTriangles = 0
	for tri in triangles:
		assert len(tri) == 3, 'length of a triangle spec is not 3, but ' + str(len(tri))
		for point in tri:
			assert point['pos'] < len(x)
			assert point['nor'] < len(nx)
			assert point['tex'] < len(u)
			daeTriangles += '%d %d %d ' % (point['pos'], point['nor'], point['tex'])
			nTriangles += 3
			
	if doubleSided:
		for tri in triangles:
			for point in tri:
				newNormal = point['nor'] + len(nx)
				newPos    = point['pos'] + len(x)
				daeTriangles += '%d %d %d ' % (newPos, newNormal, point['tex'])
				nTriangles += 3

	return daePositions, daeNormals, daeUVs, daeTriangles, nPositions, nNormals, nUVs, nTriangles


def subdivideQuadrilateralInterp(x,y,z,u,v,k):
	# use bilinear interpolation to sample the image and world coordinates.
	newX = []
	newY = []
	newZ = []
	newU = []
	newV = []
	for i in range(k+1):
		alpha = float(i) / float(k)
		worldQX = (1-alpha) * x[0] + alpha * x[1]
		worldQY = (1-alpha) * y[0] + alpha * y[1]
		worldQZ = (1-alpha) * z[0] + alpha * z[1]
		imageQU = (1-alpha) * u[0] + alpha * u[1]
		imageQV = (1-alpha) * v[0] + alpha * v[1]
		
		worldRX = (1-alpha) * x[3] + alpha * x[2]
		worldRY = (1-alpha) * y[3] + alpha * y[2]
		worldRZ = (1-alpha) * z[3] + alpha * z[2]
		imageRU = (1-alpha) * u[3] + alpha * u[2]
		imageRV = (1-alpha) * v[3] + alpha * v[2]
		for j in range(k+1):
			beta = float(j) / float(k)
			worldPX = (1-beta) * worldQX + beta * worldRX
			worldPY = (1-beta) * worldQY + beta * worldRY
			worldPZ = (1-beta) * worldQZ + beta * worldRZ
			imagePU = (1-beta) * imageQU + beta * imageRU
			imagePV = (1-beta) * imageQV + beta * imageRV
			
			newX.append(worldPX)
			newY.append(worldPY)
			newZ.append(worldPZ)
			newU.append(imagePU)
			newV.append(imagePV)
			
	# triangulate.
	endPlusOne = k*k + k - 1
	triangles = []
	nTriangles = 0
	for upperLeft in range(endPlusOne):
		if (upperLeft - k) % (k+1) == 0:
			continue
		tri = [ [], [], [] ]
		tri[0] = {'pos': upperLeft,         'tex': upperLeft,         'nor':nTriangles}
		tri[1] = {'pos': upperLeft + k + 1, 'tex': upperLeft + k + 1, 'nor':nTriangles}
		tri[2] = {'pos': upperLeft + 1,     'tex': upperLeft + 1,     'nor':nTriangles}
		triangles.append(tri);
		nTriangles = nTriangles + 1
		
		tri = [ [], [], [] ]
		tri[0] = {'pos': upperLeft + k + 1, 'tex': upperLeft + k + 1, 'nor':nTriangles}
		tri[1] = {'pos': upperLeft + k + 2, 'tex': upperLeft + k + 2, 'nor':nTriangles}
		tri[2] = {'pos': upperLeft + 1,     'tex': upperLeft + 1,     'nor':nTriangles}
		triangles.append(tri);
		nTriangles = nTriangles + 1
	
	# find the normals.
	nx,ny,nz = findNormals(newX,newY,newZ,triangles)
	
	return newX,newY,newZ,newU,newV,nx,ny,nz,triangles
# takes in
#    x,y,z as a list of positions in 3D space.
#    u,v   as a list of positions in image space.
#    k     as the number of subdivisions.
# returns
#    x,y,z as a (bigger) list of positions in 3D space.
#    u,v   as a (bigger) list of positions in image space.
#    nx,ny,nz as the normal to the plane.
#    triangles, an indexing of the triangles in the polygon.

# finds the homography between image space and the quadrilateral space.
# pros: 100% perspective correct.
# cons: assumes the input is planar.
def subdivideQuadrilateralHomography(x,y,z,u,v,k):
	assert len(x) == len(y) == len(z) == len(u) == len(v) == 4
	
	# create the axis for the plane.
	[nx,ny,nz] = normal(x[0],y[0],z[0],x[1],y[1],z[1],x[2],y[2],z[2])
	xAxis = [x[1] - x[0], y[1] - y[0], z[1] - z[0]]
	magnitude = math.sqrt(xAxis[0]*xAxis[0] + xAxis[1]*xAxis[1] + xAxis[2]*xAxis[2])
	xAxis[0] = xAxis[0] / magnitude
	xAxis[1] = xAxis[1] / magnitude
	xAxis[2] = xAxis[2] / magnitude
	
	[a,b,c] = cross(nx,ny,nz,xAxis[0],xAxis[1],xAxis[2])
	magnitude = math.sqrt(a*a + b*b + c*c)
	yAxis = [a / magnitude, b / magnitude, c / magnitude]
	
	# find out the plane coordinates of the UV points.
	qx,qy = [], []
	for i in range(len(u)):
		xDot = (x[i] - x[0]) * xAxis[0] + (y[i] - y[0]) * xAxis[1] + (z[i] - z[0]) * xAxis[2]
		yDot = (x[i] - x[0]) * yAxis[0] + (y[i] - y[0]) * yAxis[1] + (z[i] - z[0]) * yAxis[2]
		qx.append(xDot)
		qy.append(yDot)
		
	# solve for the homography between image coordinates and plane coordinates.
	A = []
	b = []
	for i in range(len(u)):
		A.append([-u[i], -v[i], -1, 0, 0, 0, qx[i] * u[i], qx[i] * v[i]])
		A.append([0, 0, 0, -u[i], -v[i], -1, qy[i] * u[i], qy[i] * v[i]])
		b.append([-qx[i]])
		b.append([-qy[i]])
	
	A = Matrix(A)
	b = Matrix(b)
	h = A.solve(b)
	H = Matrix([[h[0], h[1], h[2]], [h[3], h[4], h[5]], [h[6], h[7], 1]])
	
	# use bilinear interpolation to sample the image coordinates.
	newX = []
	newY = []
	newZ = []
	newU = []
	newV = []
	for i in range(k+1):
		alpha = float(i) / float(k)
		imageQU = (1-alpha) * u[0] + alpha * u[1]
		imageRU = (1-alpha) * u[3] + alpha * u[2]
		imageQV = (1-alpha) * v[0] + alpha * v[1]
		imageRV = (1-alpha) * v[3] + alpha * v[2]
		for j in range(k+1):
			beta = float(j) / float(k)
			imagePU = (1-beta) * imageQU + beta * imageRU
			imagePV = (1-beta) * imageQV + beta * imageRV
			
			temp = H.mmul(Matrix([[imagePU], [imagePV], [1]]))
			planeX = float(temp[0][0][0] / temp[2][0][0])
			planeY = float(temp[1][0][0] / temp[2][0][0])
			
			worldX = x[0] + planeX * xAxis[0] + planeY * yAxis[0]
			worldY = y[0] + planeX * xAxis[1] + planeY * yAxis[1]
			worldZ = z[0] + planeX * xAxis[2] + planeY * yAxis[2]
			
			newX.append(worldX)
			newY.append(worldY)
			newZ.append(worldZ)
			newU.append(imagePU)
			newV.append(imagePV)
			
	
	# make all of the triangles.
	# a little bit of reasoning, if Austin looks at this later and is confused.
	# 0      k+1     2k+2                        k(k+1) = k^2 + k
	# 1      k+2     2k+3                        k^2 + k + 1
	# 2      k+3     2k+4                        k^2 + k + 2
	# ...    ...     ...    ...    ...
	# k-1    2k      3k+1        k^2 + k - 2     k^2 + 2k - 1
	# k      2k+1    3k+2        k^2 + k - 1     k^2 + 2k
	#
	# so, iterate the upper-left side of the quadrilateral from 0 to k^2 + k - 2.
	# Then, make sure not to make any triangles on the bottom row.
	# Here are the bottom row indices:  (k, 2k+1, 3k+2, ... k^2 + k - 1, k^2 + 2k).
	# Here are the bottom row indices - k: (0, k+1, 2k+2, ... k^2 - 1, k^2 + k)
	# So, we're on the bottom row when our index mod (k+1) is 0.
	endPlusOne = k*k + k - 1
	triangles = []
	for upperLeft in range(endPlusOne):
		if (upperLeft - k) % (k+1) == 0:
			continue
		tri = [ [], [], [] ]
		tri[0] = {'pos': upperLeft,         'tex': upperLeft,         'nor':0}
		tri[1] = {'pos': upperLeft + k + 1, 'tex': upperLeft + k + 1, 'nor':0}
		tri[2] = {'pos': upperLeft + 1,     'tex': upperLeft + 1,     'nor':0}
		triangles.append(tri);
		
		tri = [ [], [], [] ]
		tri[0] = {'pos': upperLeft + k + 1, 'tex': upperLeft + k + 1, 'nor':0}
		tri[1] = {'pos': upperLeft + k + 2, 'tex': upperLeft + k + 2, 'nor':0}
		tri[2] = {'pos': upperLeft + 1,     'tex': upperLeft + 1,     'nor':0}
		triangles.append(tri);
	
	return newX,newY,newZ,newU,newV,[nx],[ny],[nz],triangles
		
	
